import os
from sklearn.model_selection import train_test_split  # conda install scikit-learn

import numpy as np
import shutil
import tqdm  # conda install tqdm
import cv2  # pip install -i https://pypi.tuna.tsinghua.edu.cn/simple opencv-python


def standardize(data, epsilon=1e-8):
    # 计算均值
    mean = data.mean()
    # 计算标准差
    std = np.std(data) + epsilon
    # 计算结果
    standardized = (data - mean) / std

    return standardized


def normalize(data, epsilon=1e-8):
    min_val = data.min()
    max_val = data.max()
    normalized = (data - min_val) / (max_val - min_val + epsilon)

    return normalized


def split_and_save_data_set(root_path, train_ratio=0.8):
    all_imgs = os.listdir(root_path)
    img_list = [file for file in all_imgs if file.endswith(".jpg")]
    GT_list = [file for file in all_imgs if file.endswith(".pgm")]

    img_list.sort()
    GT_list.sort()

    train_imgs, test_imgs, train_gts, test_gts = train_test_split(
        img_list, GT_list, train_size=train_ratio, random_state=42
    )

    train_img_dir = os.path.join(root_path, 'train', 'images')
    test_img_dir = os.path.join(root_path, 'test', 'images')
    train_gt_dir = os.path.join(root_path, 'train', 'labels')
    test_gt_dir = os.path.join(root_path, 'test', 'labels')

    os.makedirs(train_img_dir, exist_ok=True)
    os.makedirs(test_img_dir, exist_ok=True)
    os.makedirs(train_gt_dir, exist_ok=True)
    os.makedirs(test_gt_dir, exist_ok=True)

    for img in train_imgs:
        shutil.move(os.path.join(root_path, img), os.path.join(train_img_dir, img))

    for img in test_imgs:
        shutil.move(os.path.join(root_path, img), os.path.join(test_img_dir, img))

    for gt in train_gts:
        shutil.move(os.path.join(root_path, gt), os.path.join(train_gt_dir, gt))

    for gt in test_gts:
        shutil.move(os.path.join(root_path, gt), os.path.join(test_gt_dir, gt))


def preprocess(where):
    ROOT_PATH = os.path.join("./dataset", where)

    # 读入原始数据和原始标签
    img_list = os.listdir(os.path.join(ROOT_PATH, "images"))
    GT_list = os.listdir(os.path.join(ROOT_PATH, "labels"))

    img_list.sort()
    GT_list.sort()

    # 创建存储结果的文件夹
    img_dir = os.path.join(ROOT_PATH, "img_preprocessed")
    GT_dir = os.path.join(ROOT_PATH, "GT_preprocessed")

    os.makedirs(img_dir, exist_ok=True)
    os.makedirs(GT_dir, exist_ok=True)

    for index in tqdm.tqdm(range(len(img_list))):
        # 读取图像和标签
        img = cv2.imread(os.path.join(ROOT_PATH, 'images', img_list[index]))
        GT = cv2.imread(os.path.join(ROOT_PATH, 'labels', GT_list[index]),
                        cv2.IMREAD_GRAYSCALE)

        # 设黑色为0， 白色为1
        GT = np.array(GT, np.float32) / 255
        GT[GT < 0.5] = 0
        GT[GT >= 0.5] = 1

        # 标准化和归一化
        std = standardize(img)
        normalized = normalize(std)

        # 存入文件夹中
        np.save(os.path.join(img_dir, img_list[index].split(".")[0] + '.npy'), normalized)
        np.save(os.path.join(GT_dir, GT_list[index].split(".")[0] + '.npy'), GT)

if __name__ == '__main__':
    split_and_save_data_set("./dataset")
    preprocess('train')
    preprocess('test')
